﻿#region License Info

//Component of Cronos Package, http://www.codeplex.com/cronos
//Copyright (C) 2009 Anthony Brockwell

//This program is free software; you can redistribute it and/or
//modify it under the terms of the GNU General Public License
//as published by the Free Software Foundation; either version 2
//of the License, or (at your option) any later version.

//This program is distributed in the hope that it will be useful,
//but WITHOUT ANY WARRANTY; without even the implied warranty of
//MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//GNU General Public License for more details.

//You should have received a copy of the GNU General Public License
//along with this program; if not, write to the Free Software
//Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

#endregion

using System;
using System.Collections.Generic;
using System.Text;
using ABMath.IridiumExtensions;
using ABMath.ModelFramework;
using ABMath.ModelFramework.Data;
using ABMath.ModelFramework.Models;
using MathNet.Numerics.Distributions;
using MathNet.Numerics.LinearAlgebra;
using MathNet.Numerics.RandomSources;

namespace ExamplePlugin
{
    [Serializable]
    public class AdaptiveARModel : UnivariateTimeSeriesModel, IMLEEstimable
    {
        public AdaptiveARModel()
        {
            LocalInitializeParameters();
        }

        [NonSerialized]
        private TimeSeries oneStepPredAtAvail;

        private const double BandwidthCeiling = 200.0;
        private const double AdaptiveWeight = 0.4;

        /// <summary>
        /// This function provides a long description of the model, including parameter values.
        /// It is displayed in the text window when the model node is selected.
        /// </summary>
        public override string Description
        {
            get
            {
                var sb = new StringBuilder(1024);
                sb.AppendLine("Adaptive AR Model:");
                sb.AppendFormat("Mu = {0:0.0000}{1}Sigma = {2:0.0000}{3}Bandwidth = {4:0.00}{5}Average Phi = {6:0.0000}{7}",
                                Mu, Environment.NewLine, Sigma, Environment.NewLine, Bandwidth, Environment.NewLine,
                                AveragePhi, Environment.NewLine);
                return sb.ToString();
            }
        }


        private double Mu
        {
            get { return Parameters[0]; }
            set { Parameters[0] = value; }
        }

        private double Sigma
        {
            get { return Parameters[1]; }
            set { Parameters[1] = value; }
        }

        private double Bandwidth
        {
            get { return Parameters[2]; }
            set { Parameters[2] = value; }
        }

        private double AveragePhi
        {
            get { return Parameters[3]; }
            set { Parameters[3] = value;}
        }

        /// <summary>
        /// Here we must provide a way to get and set parameters for the model.
        /// </summary>
        public override Vector Parameters { get; set; }

        /// <summary>
        /// When we implement the IMLEEstimable interface, it means that users can fit models by
        /// maximum likelihood.  However, in order to do so, we have to have a function and its inverse,
        /// mapping the space of all valid parameters onto a hypercube and vice versa.
        /// </summary>
        /// <param name="param"></param>
        /// <returns></returns>
        #region IMLEEstimable Members

        public Vector ParameterToCube(Vector param)
        {
            var cube = new Vector(4);

            cube[0] = Math.Exp(param[0])/(1 + Math.Exp(param[0])); // real to [0,1]
            cube[1] = param[1]/(1 + param[1]); // real+ to [0,1]
            cube[2] = param[2]/BandwidthCeiling; // real+ to [0,1]
            cube[3] = (param[3] + 1.0)/2.0;   // [-1,1] to [0,1]

            return cube;
        }

        public Vector CubeToParameter(Vector cube)
        {
            var param = new Vector(4);

            param[0] = Math.Log(cube[0]/(1 - cube[0]));
            param[1] = cube[1]/(1.0 - cube[1]);
            param[2] = cube[2]*BandwidthCeiling;
            param[3] = cube[3]*2.0 - 1;

            return param;
        }

        #endregion

        /// <summary>
        /// This function is not required, but is useful.
        /// It provides a short string to be displayed in the model node box.
        /// </summary>
        /// <returns></returns>
        public override string GetShortDescription()
        {
            return string.Format("Adaptive AR");
        }

        /// <summary>
        /// This function returns names of individual Components of the parameter vector.
        /// They are used in dialog boxes shown when users edit parameters manually.
        /// </summary>
        /// <param name="index"></param>
        /// <returns></returns>
        public override string GetParameterName(int index)
        {
            switch (index)
            {
                case 0:
                    return "Mean";
                case 1:
                    return "Sigma";
                case 2:
                    return "Bandwidth";
                case 3:
                    return "Ave. Phi";
                default:
                    throw new ArgumentException("Invalid index");
            }
        }

        public override string GetParameterDescription(int index)
        {
            switch (index)
            {
                case 0:
                    return "Mean of the process";
                case 1:
                    return "Standard deviation of the process";
                case 2:
                    return "Number of data points used from history to generate one-step predictor";
                case 3:
                    return "Long term average of autoregressive coefficient.";
                default:
                    throw new ArgumentException("Invalid index");
            }
        }

        /// <summary>
        /// This function is called after input data and/or parameters have been changed.
        /// All outputs must be updated, including the inherited property Residuals which should
        /// be filled in with standardized residuals.
        /// GoodnessOfFit should also be set.  Usually this is just the log-likelihood.
        /// Any other model outputs should be computed here as well.
        /// </summary>
        private void ComputeResidualsAndOutputs()
        {
            var resids = new TimeSeries();
            oneStepPredAtAvail = new TimeSeries();
            double ll = 0;

            double lastPredictor = Mu;
            for (int t = 0; t < values.Count; ++t)
            {
                double predicted = GetOneStepPredictor(t, values);
                double residual = values[t] - predicted;
                ll += NormalDistributionExtensions.LogProbabilityDensity(residual, 0, Sigma);
                resids.Add(values.TimeStamp(t), (residual - Mu)/Sigma, false);
                oneStepPredAtAvail.Add(values.TimeStamp(t), lastPredictor, true);
                lastPredictor = predicted;
            }

            Residuals = resids;
            GoodnessOfFit = ll;
        }

        public override object SimulateData(object inputs, int randomSeed)
        {
            var times = inputs as List<DateTime>;
            if (times == null)
                return null; // inputs should be a list of DateTimes

            var rs = new AdditiveLaggedFibonacciRandomSource(randomSeed);
            var sd = new StandardDistribution(rs);

            var retval = new TimeSeries();
            for (int t=0 ; t<times.Count ; ++t)
            {
                double y = GetOneStepPredictor(t, retval);
                y += sd.NextDouble()*Sigma;
                retval.Add(times[t], y, false);
            }

            retval.Title = "Simul.";
            retval.Description = "Simulation from " + GetDescription();
            return retval;
        }

        public override int NumOutputs()
        {
            return base.NumOutputs() + 1;
        }

        public override object GetOutput(int socket)
        {
            int n = base.NumOutputs();
            if (socket < n)
                return base.GetOutput(socket);
            if (socket == n)
                return oneStepPredAtAvail;
            throw new SocketException();
        }

        public override string GetOutputName(int socket)
        {
            int n = base.NumOutputs();
            if (socket < n)
                return base.GetOutputName(socket);
            if (socket == n)
                return StandardOutputs.OneStepPredictorAtAvail;
            throw new SocketException();
        }

        public override object BuildForecasts(object otherData, object inputs)
        {
            throw new NotImplementedException();
        }

        private void LocalInitializeParameters()
        {
            Parameters = new Vector(4);
            Mu = 0.0;
            Sigma = 1.0;
            Bandwidth = 30;
            AveragePhi = 0;

            ParameterStates = new ParameterState[Parameters.Length];
            ParameterStates[0] = ParameterState.Consequential; // Mu follows others
            ParameterStates[1] = ParameterState.Consequential; // Sigma follows from the others
            for (int i = 2; i < Parameters.Length; ++i)
                ParameterStates[i] = ParameterState.Free; // the other coefficients are free
        }

        /// <summary>
        /// This function must fill in valid default parameters and parameter states.
        /// We call another function because the other function is also used to initialize in the constructor.
        /// </summary>
        protected override void InitializeParameters()
        {
            LocalInitializeParameters();
        }

        protected double GetOneStepPredictor(int t, TimeSeries ts)
        {
            // Here are the guts of the adaptive model.
            // start with an AR(1) with phi = 0.5 for testing
            if (t <= 3)
                return Mu;

            // carry out a regression to select autoregressive param
            int numBack = (int) Math.Ceiling(Bandwidth);
            if (numBack < 2)
                numBack = 2;
            if (numBack > t - 1)
                numBack = t - 1;
            if (numBack > BandwidthCeiling)
                numBack = (int) Math.Floor(BandwidthCeiling);

            // find \hat{\beta} using the standard formula
            double sumXX = 0;
            double sumXY = 0;
            for (int i = 0 ; i<numBack ; ++i)
            {
                double weight = 1.0;
                if (i >= (int)Math.Floor(Bandwidth))
                {
                    weight = Bandwidth - Math.Floor(Bandwidth);
                    if (weight < 0.01)
                        weight = 0.01;
                }

                double ty = ts[t - i - 1] - Mu;
                double tx = ts[t - i - 2] - Mu;
                sumXX += weight*tx*tx;
                sumXY += weight*tx*ty;
            }

            double phi = AdaptiveWeight * sumXY/sumXX + (1.0-AdaptiveWeight) * AveragePhi;

            if (phi > 1)
                phi = 1;
            if (phi < -1)
                phi = -1;

            return Mu + (ts[t - 1] - Mu)*phi;
        }

        protected override Vector ComputeConsequentialParameters(Vector parameter)
        {
            Vector newparms = parameter.Clone();
            if (ParameterStates[0] == ParameterState.Consequential)
                newparms[0] = values.SampleMean();
            if (ParameterStates[1] == ParameterState.Consequential)
            {
                // we want std. dev. of residuals
                Vector pbak = Parameters;

                double ss = 0.0;
                Parameters = parameter;
                for (int t = 0; t < values.Count; ++t)
                {
                    double residual = values[t] - GetOneStepPredictor(t, values);
                    ss += residual*residual;
                }

                Parameters = pbak;

                newparms[1] = Math.Sqrt(ss/(values.Count - 1));
            }
            return newparms;
        }

        public void CarryOutPreMLEComputations()
        {
            // nothing to do here: if there are some aspects of likelihood computation that can be reused with different parameters and the same data,
            // we would do that here
        }

        public override double LogLikelihood(Vector parameter, double penaltyFactor, bool fillOutputs)
        {
            double ll = 0;
            Vector pbak = Parameters;

            if (parameter != null)
               Parameters = parameter;

            for (int t = 0; t < values.Count; ++t)
            {
                double predicted = GetOneStepPredictor(t, values);
                double residual = values[t] - predicted;
                ll += NormalDistributionExtensions.LogProbabilityDensity(residual, 0, Sigma);
            }
            if (fillOutputs)
                ComputeResidualsAndOutputs();

            if (parameter != null)
               Parameters = pbak;
            return ll;
        }

        protected override bool CheckParameterValidity(Vector param)
        {
            // no constraints on mean
            // sigma and bandwidth must be positive
            if (Sigma < 0)
                return false;
            if (Bandwidth < 0)
                return false;
            return true;
        }

        public override Vector ComputeACF(int maxLag, bool normalize)
        {
            // for now just return white noise ACF
            var acf = new Vector(maxLag + 1);
            acf[0] = 1.0;
            return acf;
        }
    }
}